import os

import SAA as SegmentAnyAnomaly
import cv2
from utils.training_utils import *


def imgOP(file):
    # TODO 处理图片
    print(file.name)
    gpu_id = 0

    os.environ['CURL_CA_BUNDLE'] = ''
    os.environ['CUDA_VISIBLE_DEVICES'] = f"{gpu_id}"

    dino_config_file = 'GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py'
    dino_checkpoint = 'weights/groundingdino_swint_ogc.pth'
    sam_checkpoint = 'weights/sam_vit_h_4b8939.pth'
    box_threshold = 0.1
    text_threshold = 0.1
    eval_resolution = 1024
    device = f"cuda:0"
    root_dir = 'result'

    # get the model
    model = SegmentAnyAnomaly.Model(
        dino_config_file=dino_config_file,
        dino_checkpoint=dino_checkpoint,
        sam_checkpoint=sam_checkpoint,
        box_threshold=box_threshold,
        text_threshold=text_threshold,
        out_size=eval_resolution,
        device=device,
    )
    # TODO 修改模型的prompts
    textual_prompts = ['product defect. scratch. damaged. ',
                       'defect']  # detect prompts, filtered phrase
    property_text_prompts = 'the image of product defect '

    model.set_ensemble_text_prompts(textual_prompts, verbose=False)
    model.set_property_text_prompts(property_text_prompts, verbose=False)

    model = model.to(device)

    image = file
    score, appendix = model(image)

    similarity_map = appendix['similarity_map']

    image_show = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image_show = cv2.resize(image_show, (eval_resolution, eval_resolution))
    similarity_map = cv2.resize(similarity_map, (eval_resolution, eval_resolution))
    score = cv2.resize(score, (eval_resolution, eval_resolution))
    return image_show
